#include "kpextOpImgF32C2FFT.h"

#include <cmath>

using namespace kpext;


OpImgF32C2FFT::OpImgF32C2FFT(kpextVkContext *ctx) :
    BaseCompute(ctx)
{
    if(!m_vkCtx){
        throw std::runtime_error(__FUNCTION__ " construct with null pointer");
    }

    m_cols = 1;
    m_rows = 1;
    m_inBuf = nullptr;
    m_outBuf = nullptr;
    m_WkNBuf = nullptr;
    m_outputBuf = nullptr;

    m_postBarrier = true;

    //  初始化描述符集及布局
    m_descrPool = m_vkCtx->descriptorPool(
                {2, 1},
                1);

    vk::DescriptorSetLayoutBinding binding(
                0,
                vk::DescriptorType::eStorageBuffer,
                2,
                vk::ShaderStageFlagBits::eCompute
                );
    std::vector<vk::DescriptorSetLayoutBinding> bindings;
    bindings.push_back(
                binding
                );


    binding.setBinding(1);
    binding.setDescriptorType(vk::DescriptorType::eUniformTexelBuffer);
    binding.setDescriptorCount(1);
    bindings.push_back(binding);

    m_descrSetlayout = m_descrPool->createDescriptorSetLayout(
                bindings
                );

    m_descrSet = m_descrPool->recreateDescSet(*m_descrSetlayout,
                                              m_descrSet
                                              );


    // 初始化pipeline
    m_pipeline = m_vkCtx->computePipeline();
    m_pipeline->m_specConsts = std::make_shared<PipelineConstSZ>(3, 4);
    m_pipeline->m_pushConsts = std::make_shared<PipelineConstSZ>(sizeof (PushConstSZ), 1);

    m_workGroupSize = ctx->getDefWorkGroupSize();

    m_pipeline->m_specConsts->set({m_workGroupSize[0],
                                   m_workGroupSize[1],
                                   1
                                  });
    {
        auto pushconsts = reinterpret_cast<PushConstSZ *>(
                    m_pipeline->m_pushConsts->data());
        pushconsts->cols = 2;
        pushconsts->rows = 2;
        pushconsts->iscoldir = 0;
        pushconsts->bitscount = 1;
        pushconsts->calcbufStarti = 0;
        pushconsts->alphaImag = 1.0f;
        pushconsts->mulConst = 1.0f;
        pushconsts->cyc = -1;
    }

    m_pipeline->m_descrSetLayouts.push_back(m_descrSetlayout);

    if(!m_pipeline->rebuild(
                sizeof (spirv::ImgF32C2FFT), spirv::ImgF32C2FFT)
            ){
        SPDLOG_WARN("Rebuild kpext::ComputePipeline failed !");
    }

    //  工作组数量
    m_workCount = calc2DWorkGroup(1, 1,
                                  m_workGroupSize);
}

OpImgF32C2FFT::~OpImgF32C2FFT()
{

}

void OpImgF32C2FFT::rebuild(int cols, int rows,
                            std::shared_ptr<BaseBuffer> inbuf,
                            std::shared_ptr<BaseBuffer> tempbuf,
                            std::shared_ptr<CommandBuffer> cmdbuf,
                            uint32_t dir,
                            float preMul,
                            bool invert)
{
    if(!inbuf || !tempbuf || !cmdbuf || !m_vkCtx){
        throw std::runtime_error(__FUNCTION__ " input with null pointer!");
    }

    if(cols < 0) cols = 2;
    if(rows < 0) rows = 2;
    if(dir > 1) dir = 0;
    int bitscount = 1;
    {
        int temp;
        if(dir == 1){
            temp = rows;
        }else{
            temp = cols;
        }
        for (int i = 25; i >= 0; --i) {
            if(temp & (1 << i)){
                bitscount = i;
                temp = (1 << i);
                break;
            }
        }
        if(dir == 1){
            if(temp != rows){
                SPDLOG_ERROR(__FUNCTION__ " rows must be a positive integer power of 2!");
            }
            rows = temp;
        }else{
            if(temp != cols){
                SPDLOG_ERROR(__FUNCTION__ " cols must be a positive integer power of 2!");
            }
            cols = temp;
        }
    }

    //  更新描述符集信息
    if(m_inBuf != inbuf.get()
            || m_outBuf != tempbuf.get()
            || m_WkNBuf == nullptr
            || static_cast<uint32_t>(cols) != m_cols
            || static_cast<uint32_t>(rows) != m_rows
            ){
        if(!m_descrPool){
            SPDLOG_WARN(__FUNCTION__ " DescriptorPool is invalid!");
            return;
        }

        if(!m_descrSetlayout){
            SPDLOG_WARN(__FUNCTION__ " DescriptorSetLayout is invalid!");
            return;
        }

        auto descrg1 = inbuf->descrInfoGetter();
        auto descrg2 = tempbuf->descrInfoGetter();
        if(!descrg1 || !descrg2){
            SPDLOG_ERROR(__FUNCTION__ " input with null pointer!");
            return;
        }

        m_descrPool->resetDescriptorPool();
        //  重新分配描述符集
        m_descrSet = m_descrPool->recreateDescSet(*m_descrSetlayout);

        //  写入描述信息
        auto bufinfo = descrg1->getDescriptorBufferInfo();
        vk::WriteDescriptorSet writeinfo(
                    *m_descrSet,
                    0, 0,
                    1,
                    vk::DescriptorType::eStorageBuffer,
                    nullptr,
                    bufinfo.get(),
                    nullptr
                    );
        m_descrPool->writeDescriptorSet(writeinfo);

        bufinfo = descrg2->getDescriptorBufferInfo();
        writeinfo.setDstArrayElement(1);
        writeinfo.setPBufferInfo(bufinfo.get());
        m_descrPool->writeDescriptorSet(writeinfo);

        m_inBuf = inbuf.get();
        m_outBuf = tempbuf.get();


        //  更新WkN旋转因子数据
        {
            int len = static_cast<int>(cols);
            if(dir == 1){
                len = static_cast<int>(rows);
            }
            len = len / 2;
#           define PI 3.14159265358979323846f

            auto stagingbuf = m_vkCtx->stagingBuf(static_cast<size_t>(len * 4 * 2));
            auto data = stagingbuf->data<float>();
            for (int k = 0; k < len; ++k) {
                float d = -2 * PI * static_cast<float>(k);
                d = d / static_cast<float>(len * 2);
                data[2 * k] = std::cosf(d);   //实部
                data[2 * k + 1] = std::sinf(d);  //虚部
            }
            //            SPDLOG_INFO("{}", stagingbuf->vector<float>());
            stagingbuf->unmap();
            m_WkNBuf = m_vkCtx->uniformTexelBuf(stagingbuf->getSize(),
                                                vk::Format::eR32G32Sfloat
                                                );

            cmdbuf = cmdbuf->allocateCmdBuf(true,
                                            vk::CommandPoolCreateFlagBits::eTransient
                                            );
            kpext::CmdBaseBufferCpy bufcpy;
            bufcpy.prepare(
                        m_WkNBuf,
                        stagingbuf
                        );
            kpext::CmdBuffersMemBarrier barrierTx;
            barrierTx.prepare(
                        m_WkNBuf,
                        vk::PipelineStageFlagBits::eTransfer,
                        vk::PipelineStageFlagBits::eComputeShader,
                        vk::AccessFlagBits::eTransferWrite,
                        vk::AccessFlagBits::eShaderRead
                        );
            cmdbuf->begin();
            cmdbuf->record(&bufcpy);
            cmdbuf->record(&barrierTx);
            cmdbuf->end();
            cmdbuf->submit();
            cmdbuf->waitExec();

            bufinfo = m_WkNBuf->descrInfoGetter()->getDescriptorBufferInfo();
            auto bufview = m_WkNBuf->descrInfoGetter()->getBufferView();
            writeinfo.setDstBinding(1);
            writeinfo.setDstArrayElement(0);
            writeinfo.setDescriptorCount(1);
            writeinfo.setDescriptorType(vk::DescriptorType::eUniformTexelBuffer);
            writeinfo.setPBufferInfo(bufinfo.get());
            writeinfo.setPTexelBufferView(bufview.get());
            m_descrPool->writeDescriptorSet(writeinfo);
        }

    }// if

    //  更新长宽
    STCAST_ASSIGN(m_cols, cols);
    STCAST_ASSIGN(m_rows, rows);

    //  更新屏障信息
    m_cmdBarrier.prepare(
                {inbuf, tempbuf,},
                vk::PipelineStageFlagBits::eComputeShader,
                vk::PipelineStageFlagBits::eComputeShader,
                vk::AccessFlagBits::eShaderWrite,
                vk::AccessFlagBits::eShaderRead
                );

    // 设定输入输出尺寸及其它常量
    {
        auto pushconsts = reinterpret_cast<PushConstSZ *>(
                    m_pipeline->m_pushConsts->data());
        pushconsts->cols = m_cols;
        pushconsts->rows = m_rows;
        pushconsts->iscoldir = dir;
        pushconsts->bitscount = bitscount;
        pushconsts->calcbufStarti = 0;
        if(invert){
            pushconsts->alphaImag = -1.0f;
        }else{
            pushconsts->alphaImag = 1.0f;
        }
        pushconsts->mulConst = preMul;
        pushconsts->cyc = -1;

        if(bitscount & 0x01){
            m_outputBuf = inbuf.get();
        }else{
            m_outputBuf = tempbuf.get();
        }
    }

    //  重新计算工作组大小
    m_workCount = calc2DWorkGroup(m_cols, m_rows,
                                  m_workGroupSize);
}

std::shared_ptr<BaseBuffer> OpImgF32C2FFT::whichOutput(std::shared_ptr<BaseBuffer> buf1,
                                                       std::shared_ptr<BaseBuffer> buf2)
{
    if(buf1 && buf1.get() == m_outputBuf){
        return buf1;
    }

    if(buf2 && buf2.get() == m_outputBuf){
        return buf2;
    }

    SPDLOG_WARN(__FUNCTION__ " return null pointer !");
    return nullptr;
}

void OpImgF32C2FFT::record(const vk::CommandBuffer &commandBuffer)
{
//    BaseCompute::record(commandBuffer);

    if(m_pipeline
            && m_pipeline->m_pushConsts
            && m_pipeline->m_layout
            && m_pipeline->m_pipeline
            ){
        auto pushconsts = reinterpret_cast<PushConstSZ *>(
                    m_pipeline->m_pushConsts->data());
        auto bitscount = pushconsts->bitscount;


        commandBuffer.bindPipeline(
                    vk::PipelineBindPoint::eCompute,
                    *m_pipeline->m_pipeline);


        commandBuffer.bindDescriptorSets(
                    vk::PipelineBindPoint::eCompute,
                    *m_pipeline->m_layout,
                    0,
                    1,
                    m_descrSet.get(),
                    0,
                    nullptr
                    );

        do {
            commandBuffer.pushConstants(
                        *m_pipeline->m_layout,
                        vk::ShaderStageFlagBits::eCompute,
                        0,
                        static_cast<uint32_t>(m_pipeline->m_pushConsts->size()),
                        m_pipeline->m_pushConsts->data()
                        );

            commandBuffer.dispatch(
                        m_workCount[0],
                    m_workCount[1],
                    m_workCount[2]
                    );

            m_cmdBarrier.record(commandBuffer);

            pushconsts->calcbufStarti = 1 - pushconsts->calcbufStarti;
            pushconsts->cyc++;
//                        break;
        } while (bitscount--);

    }else{
        SPDLOG_ERROR(__FUNCTION__ " invalid internal var!");
    }
}
